from __future__ import print_function
from __future__ import division
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Function
from torch.nn import Parameter
import math

# https://blog.csdn.net/tsq292978891/article/details/79364140


class Criterion(nn.Module):
    def __init__(self, in_features, out_features, ):
        super(Criterion, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(torch.Tensor(out_features, in_features))
        self.reset_parameters()

    def predict(self, input):
        return F.linear(input, self.weight).detach()

    def get_accuracy(self):
        def accuracy(input, targs):
            "Computes accuracy with `targs` when `input` is bs * n_classes."
            n = targs.shape[0]
            input = self.predict(input)
            input = input.argmax(dim=-1).view(n, -1)
            targs = targs.view(n, -1)
            return (input == targs).float().mean()

        return accuracy


class SoftmaxLoss(Criterion):
    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))

    def forward(self, input, label):
        return F.cross_entropy(F.linear(input, self.weight, ), label)

    def extra_repr(self):
        return 'in_features={}, out_features={},'.format(
            self.in_features, self.out_features
        )

class ArcLoss(Criterion):
    def __init__(self, in_features, out_features, s=1,m=0.2):
        super(ArcLoss, self).__init__(in_features, out_features)
        self.register_parameter('s', nn.Parameter(torch.Tensor([s])))
        self.register_buffer('m',torch.Tensor([m]))

    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))

    def forward(self, input, label):
        self.one_hot = torch.zeros((label.size(0), self.out_features), dtype=torch.uint8).cuda()
        self.one_hot.scatter_(1, label.view(-1, 1).long(), 1)

        # --------------------------- cos(theta) & phi(theta) ---------------------------
        # 由于精度问题导致越界
        cosine = F.linear(F.normalize(input), F.normalize(self.weight)).clamp_(0, 1 - 1e-5)

        theta = cosine.acos()
        theta_m = theta + self.m
        # -------------torch.where(out_i = {x_i if condition_i else y_i) -------------
        self.output = torch.where(self.one_hot, theta_m, theta)
        return F.cross_entropy(- self.output * self.s, label)

    def extra_repr(self):
        return 'in_features={}, out_features={}, s={}, m={}'.format(
            self.in_features, self.out_features,
            self.s if 's' in dir(self) else None,
            self.m if 'm' in dir(self) else None,
        )

class ArcFaceLoss(Criterion):
    r"""Implement of large margin arc distance: :
        Args:
            in_features: size of each input sample
            out_features: size of each output sample
            s: norm of input feature
            m: margin
            cos(theta + m)
        """
    def __init__(self, in_features, out_features, s=30.0, m=0.50):
        super(ArcFaceLoss, self).__init__(in_features, out_features)
        self.s = s
        self.m = m

    def reset_parameters(self):
        nn.init.xavier_uniform_(self.weight)

    def forward(self, input, label):
        self.one_hot = torch.zeros((label.size(0), self.out_features), dtype=torch.uint8).cuda()
        self.one_hot.scatter_(1, label.view(-1, 1).long(), 1)

        # --------------------------- cos(theta) & phi(theta) ---------------------------
        cosine = F.linear(F.normalize(input), F.normalize(self.weight)).clamp_(0, 1 - 1e-5)

        theta = cosine.acos()
        theta_m = theta + self.m
        # -------------torch.where(out_i = {x_i if condition_i else y_i) -------------
        self.output = torch.where(self.one_hot, theta_m.cos(), cosine)
        return F.cross_entropy(self.output * self.s, label)

    def extra_repr(self):
        return 'in_features={}, out_features={}, s={}, m={}'.format(
            self.in_features, self.out_features,
            self.s if 's' in dir(self) else None,
            self.m if 'm' in dir(self) else None,
        )

